package com.cloudea.learn.layer;

import com.cloudea.learn.Tensor;

public class ElementSum extends Layer {
    @Override
    protected Tensor forward(Context context, Tensor... inputs) {
        Tensor c = new Tensor(1);
        double sum = 0;
        for (Tensor input: inputs){
            for (int i = 0; i < input.getSize(); i++){
                sum += input.getByLocation(i);
            }
        }
        c.set(sum, 0);
        return c;
    }

    @Override
    public Tensor[] backward(Context context, Tensor grid, Tensor... inputs) {
        Tensor[] grids = new Tensor[inputs.length];
        for (int i = 0; i < inputs.length; i++){
            grids[i] = new Tensor(inputs[i].getShape());
            for (int j = 0; j < inputs[i].getSize(); j++){
                grids[i].setByLocation(j, grid.get(0));
            }
        }
        return grids;
    }
}
